/**
 * Definition for binary tree
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    void recoverTree(TreeNode *root) {
        // Start typing your C/C++ solution below
        // DO NOT write int main() function
        TreeNode *node0 = NULL;
        TreeNode *node1 = NULL;
        TreeNode *lastNode = NULL;

        dfs(root, lastNode, node0, node1);
        if (node0 && node1) {
            swap(node0->val, node1->val);
        }
    }

    void dfs(TreeNode *root, TreeNode *&lastNode,
             TreeNode *&node0, TreeNode *&node1) {
        if (root == NULL) {
            return;
        }
        if (root->left) {
            dfs(root->left, lastNode, node0, node1);
        }
        if (lastNode != NULL && lastNode->val > root->val) {
            if (node0 == NULL) {
                node0 = lastNode;
            }
            node1 = root;
        }
        lastNode = root;
        if (root->right) {
            dfs(root->right, lastNode, node0, node1);
        }
    }
};
